"""
编码工具类，支持多种主流加密方式
"""

import base64
import hashlib
from typing import Union

try:
    from Crypto.Cipher import DES3, AES
    from Crypto.Util.Padding import pad, unpad
    from Crypto.Random import get_random_bytes
    from Crypto.PublicKey import RSA
    from Crypto.Cipher import PKCS1_OAEP
    CRYPTO_AVAILABLE = True
except ImportError:
    CRYPTO_AVAILABLE = False
    print("警告: pycryptodome库未安装，部分加密功能将不可用。请运行: pip install pycryptodome")


class Encoder:
    """
    编码器类，支持多种加密方式
    """

    @staticmethod
    def base64_encode(text: str) -> str:
        """
        Base64编码

        :param text: 待编码的字符串
        :return: 编码后的字符串
        """
        return base64.b64encode(text.encode('utf-8')).decode('utf-8')

    @staticmethod
    def md5_encode(text: str) -> str:
        """
        MD5哈希

        :param text: 待哈希的字符串
        :return: 32位小写MD5值
        """
        md5 = hashlib.md5()
        md5.update(text.encode('utf-8'))
        return md5.hexdigest()

    @staticmethod
    def sha1_encode(text: str) -> str:
        """
        SHA1哈希

        :param text: 待哈希的字符串
        :return: SHA1值
        """
        sha1 = hashlib.sha1()
        sha1.update(text.encode('utf-8'))
        return sha1.hexdigest()

    @staticmethod
    def sha256_encode(text: str) -> str:
        """
        SHA256哈希

        :param text: 待哈希的字符串
        :return: SHA256值
        """
        sha256 = hashlib.sha256()
        sha256.update(text.encode('utf-8'))
        return sha256.hexdigest()

    @staticmethod
    def des3_encrypt(text: str, key: str) -> str:
        """
        3DES加密

        :param text: 待加密的字符串
        :param key: 密钥（长度必须为24字节，或自动填充）
        :return: 加密后的base64编码字符串
        """
        if not CRYPTO_AVAILABLE:
            raise Exception("加密库不可用，请安装pycryptodome")

        # 调整密钥长度为24字节
        if len(key) < 24:
            key = key.ljust(24, '\0')
        elif len(key) > 24:
            key = key[:24]

        # 创建3DES对象
        cipher = DES3.new(key.encode('utf-8'), DES3.MODE_ECB)
        
        # 填充明文
        padded_text = pad(text.encode('utf-8'), DES3.block_size)
        
        # 加密
        encrypted = cipher.encrypt(padded_text)
        
        # 返回base64编码的结果
        return base64.b64encode(encrypted).decode('utf-8')

    @staticmethod
    def aes_encrypt(text: str, key: str) -> str:
        """
        AES加密

        :param text: 待加密的字符串
        :param key: 密钥（长度必须为16、24或32字节）
        :return: 加密后的base64编码字符串
        """
        if not CRYPTO_AVAILABLE:
            raise Exception("加密库不可用，请安装pycryptodome")

        # 调整密钥长度
        key_bytes = len(key.encode('utf-8'))
        if key_bytes < 16:
            key = key.ljust(16, '\0')
        elif key_bytes > 32:
            key = key[:32]
        elif 16 < key_bytes < 24:
            key = key[:16]
        elif 24 < key_bytes < 32:
            key = key[:24]

        # 转换为字节
        key = key.encode('utf-8')
        
        # 创建AES对象
        cipher = AES.new(key, AES.MODE_ECB)
        
        # 填充明文
        padded_text = pad(text.encode('utf-8'), AES.block_size)
        
        # 加密
        encrypted = cipher.encrypt(padded_text)
        
        # 返回base64编码的结果
        return base64.b64encode(encrypted).decode('utf-8')

    @staticmethod
    def rsa_encrypt(text: str, public_key: Union[str, bytes]) -> str:
        """
        RSA加密

        :param text: 待加密的字符串
        :param public_key: 公钥路径或公钥内容
        :return: 加密后的base64编码字符串
        """
        if not CRYPTO_AVAILABLE:
            raise Exception("加密库不可用，请安装pycryptodome")

        # 判断是路径还是密钥内容
        if public_key.startswith("-----BEGIN PUBLIC KEY-----"):
            # 是密钥内容
            rsa_key = RSA.import_key(public_key)
        else:
            # 是路径
            with open(public_key, 'r') as f:
                rsa_key = RSA.import_key(f.read())

        cipher = PKCS1_OAEP.new(rsa_key)
        encrypted = cipher.encrypt(text.encode('utf-8'))
        return base64.b64encode(encrypted).decode('utf-8')